# 导入必要的库
import pulp  # 用于线性规划和优化模型的库
import pandas as pd  # 用于数据处理和读取Excel文件
from openpyxl import load_workbook  # 用于读取和操作Excel工作簿的库

# 读取用户上传的数据文件
file_path = '附件1.xlsx'  # 包含耕地信息的文件
file_path_2 = '附件2.xlsx'  # 包含作物数据的文件
file_path_3 = '附件3.xlsx'  # 包含作物销售数据的文件
file_path_4 = 'result1_1.xlsx'  # 结果输出文件路径

# 加载Excel工作簿，获取“乡村的现有耕地”表单的内容
book = load_workbook(file_path)  # 使用openpyxl库加载指定的Excel文件
df_land = pd.read_excel(file_path, sheet_name='乡村的现有耕地')  # 读取耕地信息表
df_crops = pd.read_excel(file_path_2, sheet_name='2023年统计的相关数据')  # 读取2023年作物统计数据表
df_sells = pd.read_excel(file_path_3)  # 读取销售数据表

# 定义豆类作物的编号集合
# 作物编号对应的是Excel中的作物标识符，豆类作物编号为 1-5
bean_crops = {1, 2, 3, 4, 5}  # 用于单季作物中的豆类编号集合
# 豆类蔬菜编号为 17-19，适用于双季地块
bean_vegs = {17, 18, 19}  # 用于双季作物中的豆类蔬菜编号集合

# 整理地块类型并分类
# A 是一个字典，用于存储地块的名称（键）与地块的面积（值，单位为亩）
# 遍历df_land中的每一行，使用地块名称作为键，地块面积作为值
A = {row['地块名称']: row['地块面积/亩'] for _, row in df_land.iterrows()}

# 计算预期销售量
# 预期销售量的计算公式是：预期销售量 = 地块面积 * 亩产量
# expected_sales 是一个字典，键是作物编号，值是产量合计
expected_sales = {row['作物编号']: row['产量合计'] for _, row in df_sells.iterrows()}

# 分类单季地块
# 单季地块的类型包括：平旱地、梯田、山坡地
# plots_single_season 是一个列表，存储了所有符合条件的单季地块的名称
plots_single_season = df_land[df_land['地块类型'].isin(['平旱地', '梯田', '山坡地'])]['地块名称'].tolist()


# 双季地块：水浇地、普通大棚、智慧大棚
plots_double_season = df_land[df_land['地块类型'].isin(['水浇地', '普通大棚', '智慧大棚'])]['地块名称'].tolist()
#将双季节再进行细分一下
irrigated_land = df_land[df_land['地块类型'] == '水浇地']['地块名称'].tolist()
greenhouse_land = df_land[df_land['地块类型'] == '普通大棚']['地块名称'].tolist()
smart_greenhouse_land = df_land[df_land['地块类型'] == '智慧大棚']['地块名称'].tolist()
# 创建问题实例
model = pulp.LpProblem("Crop_Planting_Optimization", pulp.LpMaximize)

# 定义作物类型
crops_single_season = list(range(1, 16))  # 作物1-15，适用于单季地块
crops_double_season_irrigated = list(range(16, 38))  # 作物16-37，适用于水浇地
crops_double_season_greenhouse = list(range(23, 42))  # 作物23-41，适用于普通大棚
crops_double_season_smart = list(range(27, 35))  # 作物27-34，适用于智慧大棚

# 处理销售单价列，取范围的平均值
df_crops['销售单价(元/斤)'] = df_crops['销售单价/(元/斤)'].apply(
    lambda x: (float(x.split('-')[0]) + float(x.split('-')[1])) / 2 if isinstance(x, str) else x
)

# 生成销售单价、亩产量和种植成本的字典
P = {row['作物编号']: row['销售单价(元/斤)'] for _, row in df_crops.iterrows()}  # 销售单价字典
Y = {row['作物编号']: row['亩产量/斤'] for _, row in df_crops.iterrows()}  # 亩产量字典
C = {row['作物编号']: row['种植成本/(元/亩)'] for _, row in df_crops.iterrows()}  # 种植成本字典

# 定义决策变量 x[i][k][t] 是地块 i 在年份 t 的季节 j 是否种植作物 k（取值为0, 0.5或1）、
years = list(range(2024, 2031))  # 从2024到2030
x = pulp.LpVariable.dicts("x", (plots_single_season + plots_double_season,
                                crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                years, [1, 2]), 0, 1, cat='Continuous')
# 定义辅助二进制变量 y1 和 y2
y1 = pulp.LpVariable.dicts("y1", (plots_single_season + plots_double_season,
                                  crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                  years, [1, 2]), 0, 1, cat='Binary')

y2 = pulp.LpVariable.dicts("y2", (plots_single_season + plots_double_season,
                                  crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                  years, [1, 2]), 0, 1, cat='Binary')
#新变量：超出部分的产量
excess = pulp.LpVariable.dicts("excess", (plots_single_season + irrigated_land + greenhouse_land + smart_greenhouse_land,
                                          crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart,
                                          years, [1, 2]), lowBound=0, cat='Continuous')

# 目标函数：最大化所有地块的净收益
# 单季地块目标函数
# 单季地块目标函数
# 单季地块目标函数
Z_single = pulp.lpSum(
    P[k] * min(Y[k] * A[i], expected_sales[k]) * x[i][k][t][1]  # 正常销售部分
    # + 0.5 * P[k] * max(0, Y[k] * A[i] - expected_sales[k]) * x[i][k][t][1]  # 滞销部分处理，相较于第二问直接quit
    - C[k] * A[i] * x[i][k][t][1]  # 种植成本
    for t in years
    for i in plots_single_season
    for k in crops_single_season
)

# 双季地块目标函数
Z_irrigated = pulp.lpSum(
    P[k] * min(Y[k] * A[i], expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 正常销售部分
    # + 0.5 * P[k] * max(0, Y[k] * A[i] - expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 滞销部分处理
    - C[k] * A[i] * (x[i][k][t][1] + x[i][k][t][2])  # 种植成本
    for t in years
    for i in irrigated_land
    for k in crops_double_season_irrigated
)

# 普通大棚目标函数
Z_greenhouse = pulp.lpSum(
    P[k] * min(Y[k] * A[i], expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 正常销售部分
    # + 0.5 * P[k] * max(0, Y[k] * A[i] - expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 滞销部分处理
    - C[k] * A[i] * (x[i][k][t][1] + x[i][k][t][2])  # 种植成本
    for t in years
    for i in greenhouse_land
    for k in crops_double_season_greenhouse
)

# 智慧大棚目标函数
Z_smart_greenhouse = pulp.lpSum(
    P[k] * min(Y[k] * A[i], expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 正常销售部分
    # + 0.5 * P[k] * max(0, Y[k] * A[i] - expected_sales[k]) * (x[i][k][t][1] + x[i][k][t][2])  # 滞销部分处理
    - C[k] * A[i] * (x[i][k][t][1] + x[i][k][t][2])  # 种植成本
    for t in years
    for i in smart_greenhouse_land
    for k in crops_double_season_smart
)


# 综合目标函数
model += Z_single + Z_irrigated + Z_greenhouse + Z_smart_greenhouse


#x变量约束条件
for i in plots_single_season + irrigated_land + greenhouse_land + smart_greenhouse_land:
    for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
        for t in years:
            for j in [1, 2]:
                # 约束 x 的值等于 0.5 * y1 + y2
                model += x[i][k][t][j] == 0.5 * y1[i][k][t][j] + y2[i][k][t][j]
                # 确保 y1 和 y2 不会同时为 1
                model += y1[i][k][t][j] + y2[i][k][t][j] <= 1
# 约束条件1：为每种地块类型添加作物选择约束
for t in years:
    # 单季地块（平旱地、梯田、山坡地）限制作物选择
    for i in plots_single_season:
        # 在每年只能从作物 {1, 2, ..., 15} 中选择
        model += pulp.lpSum(x[i][k][t][1] for k in crops_single_season) == 1
    # 双季地块
    # 水浇地
    for i in irrigated_land:
        # 第一季可以在 {16, 17, ..., 34} 中选择
        model += pulp.lpSum(x[i][k][t][1] for k in range(16, 35)) == 1
        # 第二季只能在 {16, 35, 36, 37} 中选择
        model += pulp.lpSum(x[i][k][t][2] for k in [16, 35, 36, 37]) == 1

    # 普通大棚
    for i in greenhouse_land:
        # 第一季可以在 {17, 18, ..., 34} 中选择
        model += pulp.lpSum(x[i][k][t][1] for k in range(17, 35)) == 1
        # 第二季只能在 {38, 39, 40, 41} 中选择
        model += pulp.lpSum(x[i][k][t][2] for k in range(38, 42)) == 1

    # 智慧大棚
    for i in smart_greenhouse_land:
        # 第一季和第二季都只能在 {17, 18, ..., 34} 中选择
        model += pulp.lpSum(x[i][k][t][1] for k in range(17, 35)) == 1
        model += pulp.lpSum(x[i][k][t][2] for k in range(17, 35)) == 1

# 约束条件2：在每块地上每季的作物种植比例之和为1
for t in years:
    for i in plots_single_season:
        model += pulp.lpSum(x[i][k][t][1] for k in crops_single_season) == 1

    for i in plots_double_season:
        for j in [1, 2]:
            model += pulp.lpSum(x[i][k][t][j] for k in crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart) == 1
# 约束条件3：轮作约束
for t in years[:-1]:
    for i in plots_single_season:
        for k in crops_single_season:
            model += x[i][k][t][1] + x[i][k][t+1][1] <= 1

    for i in plots_double_season:
        for k in crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
            model += x[i][k][t][1] + x[i][k][t][2] <= 1
            model += x[i][k][t][2] + x[i][k][t+1][1] <= 1

# 轮作约束：每块地在每连续三年内至少种植一次豆类作物
for t_start in range(2024, 2028):
    for i in plots_single_season:
        model += pulp.lpSum(x[i][k][t][1] for k in bean_crops for t in range(t_start, t_start + 3)) >= 1

    for i in irrigated_land + greenhouse_land + smart_greenhouse_land:
        model += pulp.lpSum(x[i][k][t][1] + x[i][k][t][2] for k in bean_vegs for t in range(t_start, t_start + 3)) >= 1


model.solve()

results = []

for i in plots_single_season + plots_double_season:
    for k in crops_single_season + crops_double_season_irrigated + crops_double_season_greenhouse + crops_double_season_smart:
        for t in years:
            for j in [1, 2]:
                if x[i][k][t][j].varValue is not None and x[i][k][t][j].varValue > 0:
                    results.append([i, k, t, j, x[i][k][t][j].varValue])

#使用 pandas 将结果转换为数据框
print(results)
df_results = pd.DataFrame(results, columns=["地块", "作物", "年份", "季节", "种植比例"])

#显示结果
print(df_results)


# 读取附件数据
df_land = pd.read_excel('附件1.xlsx', sheet_name='乡村的现有耕地')
df_crops = pd.read_excel('附件2.xlsx', sheet_name='2023年统计的相关数据')
df_crop_names = pd.read_excel('附件1.xlsx', sheet_name='乡村种植的农作物')


result1_1 = pd.read_excel('result1_1.xlsx', sheet_name=None)

# 读取作物名称与编号的对应关系
crop_name_mapping = {row['作物编号']: row['作物名称'] for _, row in df_crop_names.iterrows()}
place_space_mapping = {row['地块名称']: row['地块面积/亩'] for _, row in df_land.iterrows()}

print(type(df_results))
# 遍历每年的优化结果并填入 result2.xlsx 对应的表格
for year in range(2024, 2031):
    # 读取该年份对应的表格
    sheet_name = f'{year}'
    sheet_data = result1_1[sheet_name]
    # index0 = sheet_data.index[sheet_data['地块名'] == 'A1'].tolist()
    # print(index0, type(index0))

    # 遍历 df_results 中对应年份的数据
    for _, row in df_results[df_results['年份'] == year].iterrows():
        plot_name = row['地块']  # 地块名称
        crop_id = row['作物']  # 作物编号
        season = row['季节']   # 第几季种植
        planting_ratio = row['种植比例']  # 种植比例
        # print(planting_ratio, type(planting_ratio))
        # 获取作物名称
        crop_name = crop_name_mapping.get(crop_id, None)
        if crop_name is None:
            continue
        land_space = place_space_mapping.get(plot_name, 0)
        if land_space == 0:
            continue
        index_list = sheet_data.index[sheet_data['地块名']==plot_name].tolist()
        index_row = index_list[season-1]
        # 根据季节和地块名称直接插入种植比例
        sheet_data.loc[index_row, crop_name] = planting_ratio * land_space

    # 更新 result2 中该年的数据
    result1_1[sheet_name] = sheet_data

# 保存修改后的 result2.xlsx
with pd.ExcelWriter('result1_1.xlsx', engine='openpyxl') as writer:
    for sheet_name, data in result1_1.items():
        data.to_excel(writer, sheet_name=sheet_name, index=False)

print("已将结果插入到 result1_1.xlsx")